#include "cartographer/mapping/pose_extrapolator.h"
#include <algorithm>
#include "absl/memory/memory.h"
#include "cartographer/transform/transform.h"
#include "glog/logging.h"

namespace cartographer
{
    namespace mapping
    {
        PoseExtrapolator::PoseExtrapolator(const common::Duration pose_queue_duration, double imu_gravity_time_constant)
            : pose_queue_duration_(pose_queue_duration)
            , gravity_time_constant_(imu_gravity_time_constant)
            , cached_extrapolated_pose_{common::Time::min(), transform::Rigid3d::Identity()}
        {
            // init pose integrator
            poseIntegrator_ = std::make_shared<PoseIntegrator>();
            poseIntegrator_->SetFilter(imu_gravity_time_constant);
            poseOptimizer_ = std::make_shared<ImuIntegrateOptimizer>(ImuOptions::GetInstance().optimizer_win_size);
        }

        std::unique_ptr<PoseExtrapolator> PoseExtrapolator::InitializeWithImu(const common::Duration pose_queue_duration,
                                                                              const double imu_gravity_time_constant,
                                                                              const sensor::ImuData &imu_data)
        {
            auto extrapolator = absl::make_unique<PoseExtrapolator>(pose_queue_duration, imu_gravity_time_constant);
            extrapolator->AddImuData(imu_data);
            extrapolator->imu_tracker_ = absl::make_unique<ImuTracker>(imu_gravity_time_constant, imu_data.time);
            extrapolator->imu_tracker_->AddImuLinearAccelerationObservation(imu_data.linear_acceleration);
            extrapolator->imu_tracker_->AddImuAngularVelocityObservation(imu_data.angular_velocity);
            extrapolator->imu_tracker_->Advance(imu_data.time);
            extrapolator->AddPose(imu_data.time, transform::Rigid3d::Rotation(extrapolator->imu_tracker_->orientation()));
            return extrapolator;
        }

        void PoseExtrapolator::ScanMatchSuccess(const common::Time &time)
        {
        }

        void PoseExtrapolator::AddPose(const common::Time time, const transform::Rigid3d &pose)
        {
            if (imu_tracker_ == nullptr)
            {
                common::Time tracker_start = time;
                if (!imu_data_.empty())
                {
                    tracker_start = std::min(tracker_start, imu_data_.front().time);
                }
                imu_tracker_ = absl::make_unique<ImuTracker>(gravity_time_constant_, tracker_start);
            }
            timed_pose_queue_.push_back(TimedPose{time, pose});
            // 只保持两个位姿
            TrimPose(time);
            // 根据两个位姿差除以时间得到线速
            UpdateVelocitiesFromPoses();
            Optimize();
            // 新的位姿到达，用IMU数据来推算新的姿态角;
            AdvanceImuTracker(time, imu_tracker_.get());
            // 用新的位姿时间裁剪IMU和ODOM
            TrimImuData();
            TrimOdometryData();
            // 三个ImuTracker其实都是只用IMU数据维护机器人的姿态；区别是odometry_imu_tracker_要
            // 根据odom时间推测姿态,extrapolation_imu_tracker_要根据雷达点时间推测姿态
            odometry_imu_tracker_ = absl::make_unique<ImuTracker>(*imu_tracker_);
            // 专门用于外部插值
            extrapolation_imu_tracker_ = absl::make_unique<ImuTracker>(*imu_tracker_);
        }

        void PoseExtrapolator::Optimize()
        {
            auto &inst = ImuOptions::GetInstance();
            if (inst.use_imu_integrate && inst.use_imu_optimize)
            {
                ImuOptNode resNode;
                const auto &lastPose = timed_pose_queue_.back();
                if (poseOptimizer_->Optimize(resNode) > 0)
                {
                    Print("last_bias_acc", last_bias_.acc_);
                    Print("last_bias_gyro", last_bias_.gyro_);
                    last_bias_.acc_ = resNode.Ba;
                    last_bias_.gyro_ = resNode.Bg;
                    Print("new_bias_acc", last_bias_.acc_);
                    Print("new_bias_gyro", last_bias_.gyro_);
                }
                auto optNode = TimedPoseToImuOptNode(lastPose, GetLinearVelocity(), last_bias_);
                optNode->integrator_ = std::make_shared<ImuPreintegrate>(last_bias_, optNode->time);
                poseOptimizer_->Push(optNode);
            }
        }

        void PoseExtrapolator::AddImuData(const sensor::ImuData &imu_data)
        {
            CHECK(timed_pose_queue_.empty() || imu_data.time >= timed_pose_queue_.back().time);
            // if (ImuOptions::GetInstance().use_imu_integrate && imu_data_.size() > 0)
            // {
            //     static ImuMeasure last = imu_data_.back();
            //     static ImuMeasure curr = imu_data;
            //     curr = imuFilter_->Estimate(last, curr);
            //     imu_data_.push_back(sensor::ImuData{curr.time, curr.acceleration, curr.angular});
            // }
            imu_data_.push_back(imu_data);
            TrimImuData();
        }

        void PoseExtrapolator::AddOdometryData(const sensor::OdometryData &odometry_data)
        {
            CHECK(timed_pose_queue_.empty() || odometry_data.time >= timed_pose_queue_.back().time);
            odometry_data_.push_back(odometry_data);
            TrimOdometryData();
            if (odometry_data_.size() < 2)
                return;
            const auto &odom_oldest = odometry_data_.front();
            const auto &odom_newest = odometry_data_.back();
            const double time_delta = common::ToSeconds(odom_oldest.time - odom_newest.time);
            const transform::Rigid3d pose_delta = odom_newest.pose.inverse() * odom_oldest.pose;
            angular_velocity_from_odometry_ = transform::RotationQuaternionToAngleAxisVector(pose_delta.rotation()) / time_delta;
            if (timed_pose_queue_.empty())
            {
                return;
            }
            // linear_velocity_in_tracking_frame_at_newest_odometry_time
            const auto linear_velocity = pose_delta.translation() / time_delta;
            const auto rotation_delta = ExtrapolateRotation(odom_newest.time, odometry_imu_tracker_.get());
            const auto orientation_odom = timed_pose_queue_.back().pose.rotation() * rotation_delta;
            linear_velocity_from_odometry_ = orientation_odom * linear_velocity;
        }

        void PoseExtrapolator::UpdateExtrapolatePose(const common::Time time)
        {
            const TimedPose &newest_timed_pose = timed_pose_queue_.back();
            CHECK_GE(time, newest_timed_pose.time);
            if (cached_extrapolated_pose_.time != time)
            {
                const auto translation = ExtrapolateTranslation(time) + newest_timed_pose.pose.translation();
                const auto rotation = newest_timed_pose.pose.rotation() * ExtrapolateRotation(time, extrapolation_imu_tracker_.get());
                cached_extrapolated_pose_ = TimedPose{time, transform::Rigid3d{translation, rotation}};
            }
        }

        transform::Rigid3d PoseExtrapolator::ExtrapolatePose(const common::Time time)
        {
            UpdateExtrapolatePose(time);
            return cached_extrapolated_pose_.pose;
        }

        transform::Rigid3d PoseExtrapolator::ExtrapolateRobotPose(const common::Time time)
        {
            UpdateExtrapolatePose(time);
            if (ImuOptions::GetInstance().use_imu_integrate)
                ImuIntegrate(timed_pose_queue_.back(), time);
            return cached_extrapolated_pose_.pose;
        }

        int PoseExtrapolator::ImuIntegrate(const TimedPose &lastPose, const common::Time currTime)
        {
            int poseSize = timed_pose_queue_.size();
            float duration = common::ToSeconds(currTime - lastPose.time);
            // LOG(INFO) << "[ExtraRobotPose] poseSize " << poseSize << " dt " << duration;
            if (poseSize < 2 || currTime <= lastPose.time)
                return -1;
            auto imuSize = imu_data_.size();
            if (imuSize < 2)
            {
                // LOG(INFO) << "[ExtraRobotPose] imu not enough " << imuSize;
                return -2;
            }
            float dtImuPose = common::ToSeconds(imu_data_.front().time - lastPose.time);
            // LOG(INFO) << "[ExtraRobotPose] imuSize " << imuSize << " dt " << dtImuPose;
            auto lastPoseTime = common::ToUniversal(lastPose.time);
            std::shared_ptr<ImuPreintegrate> integrator = nullptr;
            if (ImuOptions::GetInstance().use_imu_optimize)
            {
                // 要进行IMU优化，每个位姿对应一个优化节点，节点保存其对应的预积分器，节点的下一个位姿就是用这个预积分器计算的;
                // 需要注意的是推算下一个节点给到前端匹配可能会失败，此时预积分器已积分了，那么这里要判断一下，如果已积分则需要
                // 重新初始化integrator_
                auto optNode = poseOptimizer_->GetOptNode(lastPoseTime);
                if (optNode)
                {
                    optNode->SetBias(last_bias_);
                    if (optNode->integrator_ != nullptr)
                    {
                        if (optNode->integrator_->hasIntegrated)
                        {
                            optNode->integrator_->Initialize(last_bias_);
                        }
                    }
                    else
                    {
                        optNode->integrator_ = std::make_shared<ImuPreintegrate>(last_bias_, lastPoseTime);
                    }
                    integrator = optNode->integrator_;
                }
            }
            if (integrator == nullptr)
                integrator = std::make_shared<ImuPreintegrate>(last_bias_, lastPoseTime);
            ImuMeasSet imuSet;
            for (const auto &imuData : imu_data_)
                imuSet.emplace_back(imuData);
            poseIntegrator_->SetIntegrator(integrator);
            int integrateRes = poseIntegrator_->Integrate(imuSet, lastPose.time, currTime);
            if (integrateRes > 0)
            {
                auto lastVel = GetLinearVelocity();
                auto pvrLast = TimedPoseToPVR(lastPose, lastVel);
                auto pvrCurr = poseIntegrator_->Predict(pvrLast, last_bias_);
                auto currPose = PVRToTimedPose(pvrCurr);
                auto poseDelta = cached_extrapolated_pose_.pose.inverse() * currPose.pose;
                cached_extrapolated_pose_ = currPose;
                // PrintPose("ImuPose", currPose);
                // PrintPose("ExtPose", lastPose);
                // PrintPose("DeltaPose", poseDelta);
                // LOG(INFO) << "org_vel:" << VecToStr(lastVel) << "\n";
                // LOG(INFO) << "imu_vel:" << VecToStr(pvrCurr.V) << "\n";
            }
            // LOG(INFO) << "[pose] integrateRes " << integrateRes << "\n";
            return 1;
        }

        Eigen::Quaterniond PoseExtrapolator::EstimateGravityOrientation(const common::Time time)
        {
            ImuTracker imu_tracker = *imu_tracker_;
            AdvanceImuTracker(time, &imu_tracker);
            return imu_tracker.orientation();
        }

        void PoseExtrapolator::UpdateVelocitiesFromPoses()
        {
            if (timed_pose_queue_.size() < 2)
                return;
            CHECK(!timed_pose_queue_.empty());
            const TimedPose &newest_timed_pose = timed_pose_queue_.back();
            const auto newest_time = newest_timed_pose.time;
            const TimedPose &oldest_timed_pose = timed_pose_queue_.front();
            const auto oldest_time = oldest_timed_pose.time;
            const double queue_delta = common::ToSeconds(newest_time - oldest_time);
            // 位姿频率高于阈值则不更新速度
            if (queue_delta < common::ToSeconds(pose_queue_duration_))
            {
                LOG(WARNING) << "Queue too short for velocity estimation. Queue duration: " << queue_delta << " s";
                return;
            }
            const transform::Rigid3d &newest_pose = newest_timed_pose.pose;
            const transform::Rigid3d &oldest_pose = oldest_timed_pose.pose;
            linear_velocity_from_poses_ = (newest_pose.translation() - oldest_pose.translation()) / queue_delta;
            angular_velocity_from_poses_ =
                transform::RotationQuaternionToAngleAxisVector(oldest_pose.rotation().inverse() * newest_pose.rotation()) / queue_delta;
        }

        void PoseExtrapolator::TrimPose(const common::Time time)
        {
            while (timed_pose_queue_.size() > 2 && timed_pose_queue_[1].time <= time - pose_queue_duration_)
            {
                timed_pose_queue_.pop_front();
            }
        }

        void PoseExtrapolator::TrimImuData()
        {
            // 把小于最新pose的imu数据删掉
            while (imu_data_.size() > 1 && !timed_pose_queue_.empty() && imu_data_[1].time <= timed_pose_queue_.back().time)
            {
                imu_data_.pop_front();
            }
        }

        void PoseExtrapolator::TrimOdometryData()
        {
            while (odometry_data_.size() > 2 && !timed_pose_queue_.empty() && odometry_data_[1].time <= timed_pose_queue_.back().time)
            {
                odometry_data_.pop_front();
            }
        }

        void PoseExtrapolator::AdvanceImuTracker(const common::Time time, ImuTracker *const imu_tracker) const
        {
            CHECK_GE(time, imu_tracker->time());
            if (imu_data_.empty() || time < imu_data_.front().time)
            {
                // There is no IMU data until 'time', so we advance the ImuTracker and use
                // the angular velocities from poses and fake gravity to help 2D stability.
                imu_tracker->Advance(time);
                imu_tracker->AddImuLinearAccelerationObservation(Eigen::Vector3d::UnitZ());
                imu_tracker->AddImuAngularVelocityObservation(GetAngularVelocity());
                return;
            }
            if (imu_tracker->time() < imu_data_.front().time)
            {
                // Advance to the beginning of 'imu_data_'.
                imu_tracker->Advance(imu_data_.front().time);
            }
            // lower_bound用于找到imu_data_中第一个大于或等于imu_tracker->time()的元素
            auto it = std::lower_bound(imu_data_.begin(), imu_data_.end(), imu_tracker->time(),
                                       [](const sensor::ImuData &imu_data, const common::Time &time) { return imu_data.time < time; });
            while (it != imu_data_.end() && it->time < time)
            {
                imu_tracker->Advance(it->time);
                imu_tracker->AddImuLinearAccelerationObservation(it->linear_acceleration);
                imu_tracker->AddImuAngularVelocityObservation(it->angular_velocity);
                ++it;
            }
            imu_tracker->Advance(time);
        }

        /// imu_tracker_是跟踪Pose得到的上一时刻的姿态;该函数求的是姿态增量;
        Eigen::Quaterniond PoseExtrapolator::ExtrapolateRotation(const common::Time time, ImuTracker *const imu_tracker) const
        {
            CHECK_GE(time, imu_tracker->time());
            AdvanceImuTracker(time, imu_tracker);
            const Eigen::Quaterniond last_orientation = imu_tracker_->orientation();
            return last_orientation.inverse() * imu_tracker->orientation();
        }

        Eigen::Vector3d PoseExtrapolator::ExtrapolateTranslation(common::Time time)
        {
            const TimedPose &newest_timed_pose = timed_pose_queue_.back();
            const double deltaTime = common::ToSeconds(time - newest_timed_pose.time);
            return deltaTime * GetLinearVelocity();
        }

        PoseExtrapolator::ExtrapolationResult PoseExtrapolator::ExtrapolatePosesWithGravity(const std::vector<common::Time> &times)
        {
            std::vector<transform::Rigid3f> poses;
            for (auto it = times.begin(); it != std::prev(times.end()); ++it)
            {
                poses.push_back(ExtrapolatePose(*it).cast<float>());
            }
            return ExtrapolationResult{poses, ExtrapolatePose(times.back()), GetLinearVelocity(), EstimateGravityOrientation(times.back())};
        }

    } // namespace mapping
} // namespace cartographer
